[perf] Copy via memoryview in pack_into instead of torch.frombuffer - #154
Merged
Merged
Conversation
pack_into built a uint8 tensor per item with torch.frombuffer and dispatched a copy_ kernel for it. At TQ's object sizes the dispatch dominates: one key is stored per (field, sample), so items are a few hundred bytes each and the tensor construction plus kernel launch costs far more than the memcpy. Copy through memoryview slice assignment instead. Casting both sides to "B" keeps this correct whether the caller passes bytes, a memoryview or a numpy array. Also compute the size requirement from the item views already materialised for the copy, so calc_packed_size no longer walks every item a second time. Measured (Ascend NPU container, torch 2.9.0+cpu): 1024 objects x 512B: 35.0ms -> 5.5ms (6.4x) 7168 objects x 512B: 248.6ms -> 40.7ms (6.1x) 2048 objects x 16KB: 82.5ms -> 22.5ms (3.7x) In a real verl GRPO run, pack_into for a 512-sample 2-field put went 51.3ms -> 6.0ms, taking put_data from 183.9ms to 139.7ms single-node and 364.9ms to 290.4ms dual-node. Output is byte-for-byte identical to the previous implementation; both serial_utils test modules pass, and the packed buffers still round-trip through unpack_from/decode unchanged. Note both callers benefit: YuanrongStorageClient and MooncakeStore both go through batch_encode_into. The measurements above are from the yuanrong path; Mooncake was not measured.
Collaborator
|
Codex Review: Didn't find any major issues. Hooray! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
0oshowero0
approved these changes
Aug 12, 2026
CLA Signature Guide@Chase-Rong , thanks for your pull request. The following commit(s) are not associated with a signed Contributor License Agreement (CLA).
To sign CLA, click here. To check if your email is configured correctly, refer to the FAQs. Once you've signed the CLA or updating your email, please comment |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Problem
pack_intoconcatenates the msgpack fragments of one object into a contiguous buffer —a 4-byte header, an 8-byte
(offset, length)entry per fragment, then tightly packedpayload. It is a memcpy with an index table.
The old implementation did that memcpy through tensors:
At TQ's object sizes (one key per (sample, field), so a few hundred bytes each) none of
that work is the copy. Timed per step, 1024 objects of 512B, best of 5:
torch.frombuffer(item_mv)target_tensor[off:off+n].copy_(src)A 512-byte memcpy is tens of nanoseconds. Two tensor objects get constructed and one
dispatcher lookup happens per fragment, and slicing the target costs more than the copy
itself. Tensors are the wrong abstraction here — their fixed per-op cost does not
amortise over a few hundred bytes.
Change
Copy through memoryview slice assignment, which resolves to
PyBuffer_ToContiguous->memcpywith no object allocation and no dispatcher:memoryview(item).cast("B")target_mv[off:off+n] = item_mvCasting both sides to
"B"is required, not cosmetic: slice assignment demands matchingformat and itemsize, and callers pass
bytes, amemoryview, or a float32 numpy array.Normalising to unsigned bytes makes the assignment legal and correct for all three.
Second change in the same function: the old code built a
memoryviewper fragmenttwice — once inside
calc_packed_sizejust to read.nbytes, once in the copy loop.The views are now materialised once and shared, with
requiredcomputed inline.Measurements
verl GRPO on
mainat 750719e, this PR applied alone (the other two perf branchesreverted to baseline). Medians; two independent baseline runs per configuration are
shown to give the run-to-run spread.
Single-node, 512-sample put (1024 objects):
pack_into, per objectcalc_packed_sizecallscalc_packed_sizetotalput_dataDual-node, 1024-sample put (2048 objects):
The per-step micro-benchmark ratio (6.70x) and the production ratio (6.1x) agree, so the
mechanism above accounts for the result. Absolute values differ because production
fragments are larger and more numerous than the synthetic 512B case.
yuanrong RPC time is unchanged (32.1 -> 31.5ms single-node, 55.2 -> 61.6ms dual-node),
so the saving is entirely client-side CPU.
get_datais unaffected (medians within 3%across variants).
Note dual-node
put_datawall clock is not a reliable metric in this environment: anintermittent ~480ms
run_in_executorhandoff delay appears in some runs (two baselineruns measured 606ms and 1083ms with every instrumented step within 1%, the difference
landing entirely in un-instrumented time). Judge the dual-node effect by TQ-side total
or pack phase, which agree to 0.2%/0.9% between those two runs.
Side effect: GIL contention drops as well
batch_encode_intostill dispatchespack_intoacross 16 threads. Because each task nowholds the GIL for 1.20us instead of 8.07us, the queue stops backing up:
pack_intosum (incl. lock wait)The baseline's per-call wall clock inflates to 1942us under contention; here the
accumulated per-thread time equals the real work, i.e. threads are barely waiting on each
other. That is why the phase improves by more than the 6.7x of the copy alone.
Correctness
Output is byte-for-byte identical to the previous implementation. Both
tests/test_serial_utils_on_cpu.pyandtests/test_serial_utils_batch_on_cpu.pypass,and packed buffers still round-trip through
unpack_from/decodeunchanged. Thebuffer-too-small `ValueError